from Kitsune import Kitsune
import numpy as np
import time
import pandas as pd
import sys
import os
import threading
from multiprocessing import Process, Queue, Pipe, Value

##############################################################################
# Kitsune a lightweight online network intrusion detection system based on an ensemble of autoencoders (kitNET).
# For more information and citation, please see our NDSS'18 paper: Kitsune: An Ensemble of Autoencoders for Online Network Intrusion Detection

# This script demonstrates Kitsune's ability to incrementally learn, and detect anomalies in recorded a pcap of the Mirai Malware.
# The demo involves an m-by-n dataset with n=115 dimensions (features), and m=100,000 observations.
# Each observation is a snapshot of the network's state in terms of incremental damped statistics (see the NDSS paper for more details)

#The runtimes presented in the paper, are based on the C++ implimentation (roughly 100x faster than the python implimentation)
###################  Last Tested with Anaconda 3.6.3   #######################

# Load Mirai pcap (a recording of the Mirai botnet malware being activated)
# The first 70,000 observations are clean...
# print("Unzipping Sample Capture...")
# import zipfile
# with zipfile.ZipFile("mirai.zip","r") as zip_ref:
    # zip_ref.extractall()

def statistics_info():
    global timer, pps, bps
    pps_total = 0
    bps_total = 0
    for i in range(ad_number):
        pps_total += pps[i].value
        bps_total += bps[i].value
        pps[i].value = 0
        bps[i].value = 0
    print("%f pps, %f bps" % (pps_total / 1.0, bps_total / 1.0))
    timer = threading.Timer(1, statistics_info)
    timer.start()

timer = threading.Timer(1, statistics_info)
timer.start()

pps = []
bps = []
fe_number = 1
ad_number = 1
for i in range(ad_number):
    pps.append(Value('I',0))
    bps.append(Value('I',0))

def get_pcap_name(pcap_path):
    while pcap_path.find('/') != -1:
        pcap_path = pcap_path[pcap_path.find('/') + 1 :]
    pcap_name = pcap_path[: pcap_path.find('.')]
    return pcap_name

# File location
# path = "mirai.pcap" #the pcap, pcapng, or tsv file to process.
path = sys.argv[1] #the pcap, pcapng, or tsv file to process.
packet_limit = np.Inf #the number of packets to process

# KitNET params:
maxAE = 10 #maximum size for any autoencoder in the ensemble layer
# FMgrace = 5000 #the number of instances taken to learn the feature mapping (the ensemble's architecture)
# ADgrace = 50000 #the number of instances used to train the anomaly detector (ensemble itself)
FMgrace = int(sys.argv[2]) #the number of instances taken to learn the feature mapping (the ensemble's architecture)
ADgrace = int(sys.argv[3]) #the number of instances used to train the anomaly detector (ensemble itself)
target_path = sys.argv[4]

# Build Kitsune
K = Kitsune(path,packet_limit,maxAE,FMgrace,ADgrace)

print("Running Kitsune:")
RMSEs = []
queue_fe_ad = Queue()
pipe_fe_ad = []
for i in range(fe_number):
    for i in range(ad_number):
        pipe_fe_ad.append(Pipe())
queue_rmse = Queue()

def FE_module_function(pps, bps, fe_seq, fe_number, pipe_fe_ad):
    i = 0
    print('This is FE-%d. PID: %s, PPID: %s' %(fe_seq, os.getpid(),os.getppid()))
    while True:
        i += 1
        if i % 100000 == 0:
            print("%d packets have been processed by FE-%d" % (i, fe_seq))
        x, framelen = K.FE.get_next_vector()
        # queue_fe_ad.put(x)
        pipe_fe_ad[fe_seq * ad_number + i % ad_number][0].send(x)
        if len(x) == 0:
            break
        bps[i % ad_number].value += framelen

def AD_module_function(pps, bps, fe_number, ad_seq, ad_number, pipe_fe_ad, queue_rmse):
    i = 0
    print('This is AD-%d. PID: %s, PPID: %s' %(ad_seq, os.getpid(),os.getppid()))
    while True:
        # if not queue_fe_ad.empty():
            # x = queue_fe_ad.get()
            # if len(x) == 0:
                # break
            # rmse = K.AnomDetector.process(x)
            # queue_rmse.put(rmse)
            # pps.value += 1
        # else:
            # bps.value += 0
            # pps.value += 0
        i += 1
        if i % 100000 == 0:
            print("%d packets have been processed by AD-%d" % (i, ad_seq))
        x = pipe_fe_ad[(i % fe_number) * ad_number + ad_seq][1].recv()
        if len(x) == 0:
            break
        rmse = K.AnomDetector.process(x)
        queue_rmse.put(rmse)
        pps[ad_seq].value += 1


start = time.time()
# Here we process (train/execute) each individual packet.
# In this way, each observation is discarded after performing process() method.

FE_modules = []
AD_modules = []

for i in range(fe_number):
    FE_modules.append(Process(target=FE_module_function, args=(pps, bps, i, ad_number, pipe_fe_ad)))
for i in range(ad_number):
    AD_modules.append(Process(target=AD_module_function, args=(pps, bps, fe_number, i, ad_number, pipe_fe_ad, queue_rmse)))
for i in range(fe_number):
    FE_modules[i].start()
for i in range(ad_number):
    AD_modules[i].start()
for i in range(fe_number):
    FE_modules[i].join()
for i in range(ad_number):
    AD_modules[i].join()

stop = time.time()
print("Complete. Time elapsed: "+ str(stop - start))

while not queue_rmse.empty():
    rmse = queue_rmse.get()
    RMSEs.append(rmse)
RMSEs = list(reversed(RMSEs))

# Here we demonstrate how one can fit the RMSE scores to a log-normal distribution (useful for finding/setting a cutoff threshold \phi)
from scipy.stats import norm
benignSample = np.log(RMSEs[FMgrace+ADgrace+1:100000])
logProbs = norm.logsf(np.log(RMSEs), np.mean(benignSample), np.std(benignSample))

pd.DataFrame(RMSEs).to_csv(target_path + '/' + get_pcap_name(path) + "_RMSEs.csv", header=False, index=False)
pd.DataFrame(logProbs).to_csv(target_path + '/' + get_pcap_name(path) + "_logProbs.csv", header=False, index=False)

# plot the RMSE anomaly scores
print("Plotting results")
from matplotlib import pyplot as plt
plt.figure(figsize=(10,5))
fig = plt.scatter(range(FMgrace+ADgrace+1,len(RMSEs)),RMSEs[FMgrace+ADgrace+1:],s=0.1,c=logProbs[FMgrace+ADgrace+1:],cmap='RdYlGn')
plt.yscale("log")
plt.title("Anomaly Scores from Kitsune's Execution Phase")
plt.ylabel("RMSE (log scaled)")
plt.xlabel("Time elapsed [min]")
# plt.annotate('Mirai C&C channel opened [Telnet]', xy=(121662,RMSEs[121662]), xytext=(151662,1),arrowprops=dict(facecolor='black', shrink=0.05),)
# plt.annotate('Mirai Bot Activated\nMirai scans network\nfor vulnerable devices', xy=(122662,10), xytext=(122662,150),arrowprops=dict(facecolor='black', shrink=0.05),)
# plt.annotate('Mirai Bot launches DoS attack', xy=(370000,100), xytext=(390000,1000),arrowprops=dict(facecolor='black', shrink=0.05),)
figbar=plt.colorbar()
figbar.ax.set_ylabel('Log Probability\n ', rotation=270)
plt.savefig(target_path + '/' + get_pcap_name(path) + "_fig.pdf")
plt.show()
